package edu.kit.csl.pisa.models;

import java.util.Arrays;

import edu.kit.csl.pisa.ui.Configuration;

/**
 * This class is used to estimate parameters of the negative binomial 
 * distribution.
 * 
 * @see http://pic.dhe.ibm.com/infocenter/spssstat/v22r0m0/index.jsp?topic=%2Fcom.ibm.spss.statistics.algorithms%2Falg_simplan_distfit_discrete_negbin.htm
 */
abstract public class NBEstimator {
	// Stored training cases
	protected double sumW, sumWK; // Holds sum w_i and sum w_i*k_i
	protected double[] weights;
	
	// Estimated parameters
	protected int r;
	protected double p;
	protected int mu;
	
	/**
	 * Sole constructor.
	 * 
	 * @param mu initial value for r (p is initialized with 0.5 per default)
	 */
	public NBEstimator(int mu) {
		int maxTrgtSenLen = 
				Configuration.getSingleton().getInteger("maxTrgtSenLen");
		weights = new double[maxTrgtSenLen+1];
		this.mu = Math.max(1, mu);
		p = 0.5;
		r = this.mu;
	}
	
	/**
	 * Add a training case. This is used when {@link #estimate()} is called
	 * the next time.
	 * 
	 * @param weight weight of that training case
	 * @param nTrgtWords number of target language words (training data)
	 */
	public void addTrainingCase(double weight, int nTrgtWords) {
		weights[nTrgtWords] += weight;
	}
	
	/**
	 * Remove all stored training cases.
	 */
	public void clear() {
		Arrays.fill(weights, 0.0);
	}
	
	/**
	 * Estimate r and p parameters based on the training cases added so far.
	 * Afterwards, the parameters can be fetched using {@link #getP()} and
	 * {@link #getR()}.
	 */
	public void estimate() {
		// Iterate r till 2*maxTrgtSenLen, MLE for p
		int maxTrgtSenLen = 
				Configuration.getSingleton().getInteger("maxTrgtSenLen");
		// Update sumW and sumWK
		sumW = AlignmentModel.sum(weights);
		if (sumW <= 0.0) { // No data -> Set p,r to default values
			p = 0.5;
			r = mu;
			return;
		}
		sumWK = 0.0;
		for (int k = 1; k < weights.length; k++) {
			sumWK += weights[k] * k;
		}
		double bestL = Double.NEGATIVE_INFINITY;
		for (int r = 1; r <= 2*maxTrgtSenLen; r++) {
			double p = estimateP(r);
			double l = getLogLikelihood(p, r);
			if (l > bestL) {
				this.p = p;
				this.r = r;
				bestL = l;
			}
		}
	}
	
	/**
	 * Estimate the p parameter given the r parameter and the data collected
	 * so far. GoF design pattern template method.
	 * 
	 * @param r r parameter
	 * @return p parameter
	 */
	abstract protected double estimateP(double r);

	// Gets log likelihood of collected data given p,r parameters
	// Missing \sum w_i log(x_i!) term since it is constant
	protected double getLogLikelihood(double p, int r) {
		double gammaSum = 0.0;
		for (int k = 0; k < weights.length; k++) {
			gammaSum += weights[k] * logGamma(k + r);
		}
		// Missing \sum w_i log(x_i!) term since it is constant
		return gammaSum - sumW * logGamma(r)
			+  sumW * ((double) r) * Math.log(1.0-p) + Math.log(p) * sumWK;
	}
	
	// Returns log(Gamma(x))
	protected double logGamma(int x) {
		if (x - 1 < AlignmentModel.logFacFert.length) {
			return AlignmentModel.logFacFert[x-1];
		}
		double acc = 0.0;
		for (int i = 2; i <= x-1; i++) {
			acc += Math.log(i);
		}
		return acc;
	}

	/**
	 * @return the estimated r parameter
	 */
	public int getR() {
		return r;
	}

	/**
	 * @return the estimated p parameter
	 */
	public double getP() {
		return p;
	}
}
